{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# 初始化"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 基础导入"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69b28288",
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd '/home/xiaobo/Project/Trans-WSSS'\n",
    "%pwd\n",
    "import os.path as osp\n",
    "import sys\n",
    "import warnings\n",
    "from typing import Any\n",
    "\n",
    "import alchemy_cat.data.plugins.augers as au\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from PIL import Image\n",
    "from alchemy_cat.acplot import col_all\n",
    "from alchemy_cat.alg import size2HW\n",
    "from alchemy_cat.contrib.tasks.wsss.viz import viz_cam\n",
    "from alchemy_cat.contrib.voc import VOC_CLASSES\n",
    "from alchemy_cat.py_tools import ADict\n",
    "from skimage.feature import peak_local_max\n",
    "from skimage.measure import label, regionprops\n",
    "\n",
    "from libs.seeding.score import cam2score_cuda\n",
    "from utils.resize import resize_cam_cuda\n",
    "\n",
    "sys.path.append(\"others/segment_anything\")\n",
    "from segment_anything import sam_model_registry, SamPredictor"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## mask可视化函数"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29bc90d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_mask(mask, ax, random_color=False):\n",
    "    if random_color:\n",
    "        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n",
    "    else:\n",
    "        color = np.array([30/255, 144/255, 255/255, 0.6])\n",
    "    h, w = mask.shape[-2:]\n",
    "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
    "    ax.imshow(mask_image)\n",
    "\n",
    "\n",
    "def show_points(coords, labels, ax, marker_size=375):\n",
    "    pos_points = coords[labels==1]\n",
    "    neg_points = coords[labels==0]\n",
    "    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
    "    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   \n",
    "\n",
    "\n",
    "def show_box(box, ax):\n",
    "    x0, y0 = box[0], box[1]\n",
    "    w, h = box[2] - box[0], box[3] - box[1]\n",
    "    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## CAM/score可视化函数"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def viz_fg_cam(img: np.ndarray, lb: np.ndarray,\n",
    "               cam: np.ndarray | torch.Tensor,\n",
    "               suptitle: str, with_bg: bool=False,\n",
    "               input_points: list[np.ndarray]=None, input_labels: list[np.ndarray]=None):\n",
    "    if torch.is_tensor(cam):\n",
    "        cam = cam.detach().cpu().numpy()\n",
    "\n",
    "    pos_names = ['dummy'] + (['background'] if with_bg else []) + [VOC_CLASSES[1:][cls] for cls in fg_cls]\n",
    "\n",
    "    plt.figure(dpi=250)\n",
    "\n",
    "    viz_cam(fig=plt.gcf(),\n",
    "            img_id=img_id, img=img, label=lb,\n",
    "            cls_in_label=np.ones(len(pos_names), dtype=np.uint8),\n",
    "            cam=cam,\n",
    "            cls_names=pos_names,\n",
    "            get_row_col=col_all)\n",
    "\n",
    "    if input_points is not None:\n",
    "        for input_point, input_label, ax in zip(input_points, input_labels, plt.gcf().axes[1::2], strict=True):\n",
    "            show_points(input_point, input_label, ax, marker_size=75)\n",
    "\n",
    "    if suptitle:\n",
    "        plt.suptitle(suptitle, fontsize=8)\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def cal_sam_input_size(sam_transform, ori_size: Any, low_res: bool=True) -> tuple[int, int]:\n",
    "    ori_h, ori_w = size2HW(ori_size)\n",
    "    sam_input_h, sam_input_w = sam_transform.get_preprocess_shape(ori_h, ori_w, sam_transform.target_length)\n",
    "    if low_res:\n",
    "        sam_input_h, sam_input_w = sam_input_h // 4, sam_input_w // 4\n",
    "    return (sam_input_h, sam_input_w)\n",
    "\n",
    "\n",
    "def transform_img_lb(sam_transform, img: np.ndarray, lb: np.ndarray, low_res: bool=True) -> tuple[np.ndarray, np.ndarray]:\n",
    "    sam_input_size = cal_sam_input_size(sam_transform, img.shape[:2], low_res=low_res)\n",
    "    img, lb = au.scale_img_label(sam_input_size, img, lb, align_corner=False, PIL_mode=Image.BILINEAR)\n",
    "    img, lb = au.pad_img_label(img, lb, pad_img_to=sam_transform.target_length // (4 if low_res else 1),\n",
    "                               img_pad_val=0, ignore_label=255, pad_location='right-bottom')\n",
    "    return img, lb"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 全局数据"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 读取模型"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "sam_checkpoint = \"pretrains/SAM/sam_vit_h_4b8939.pth\"\n",
    "model_type = \"vit_h\"\n",
    "\n",
    "device = \"cuda:0\"\n",
    "\n",
    "sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\n",
    "sam.to(device=device)\n",
    "\n",
    "predictor = SamPredictor(sam)\n",
    "\n",
    "sam_transform = predictor.transform\n",
    "\n",
    "sam_img_size = predictor.model.image_encoder.img_size\n",
    "low_res_mask_size = sam_img_size // 4\n",
    "\n",
    "print(f'sam_img_size: {sam_img_size}, low_res_mask_size: {low_res_mask_size}')"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 读取数据"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 读取图片"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# * 读取和显示图片。\n",
    "img_id = '2007_000170'\n",
    "# img_id = '2007_003205'\n",
    "img = cv2.imread(f'datasets/VOC2012/JPEGImages/{img_id}.jpg')\n",
    "img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "lb = np.asarray(Image.open(f'datasets/VOC2012/SegmentationClass/{img_id}.png'), dtype=np.uint8)\n",
    "\n",
    "ori_h, ori_w = img.shape[:2]\n",
    "\n",
    "sam_low_res_input_h, sam_low_res_input_w = cal_sam_input_size(sam_transform, img.shape[:2], low_res=True)\n",
    "low_res_img, low_res_lb = transform_img_lb(sam_transform, img, lb, low_res=True)\n",
    "\n",
    "low_res_pad_h = low_res_mask_size - sam_low_res_input_h\n",
    "low_res_pad_w = low_res_mask_size - sam_low_res_input_w\n",
    "\n",
    "plt.imshow(img)\n",
    "plt.show()\n",
    "\n",
    "print(f\"ori_h: {ori_h}, ori_w: {ori_w}, \"\n",
    "      f\"sam_low_res_input_h: {sam_low_res_input_h}, sam_low_res_input_w: {sam_low_res_input_w}, \"\n",
    "      f\"low_res_pad_h: {low_res_pad_h}, low_res_pad_w: {low_res_pad_w}\")"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 读取CAM、score、affed cam、affed score、CRF后score"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# * 设置基础实验。\n",
    "infer_rslt_dir = 'experiment/clip_cam/离线伪真,CI/l1/ps自cl_loss,5100/infer/final'\n",
    "aff_rslt_dir = osp.join(infer_rslt_dir, 'aff2次,at_cam,att1次,·5掩阈,ce_npp,mask')\n",
    "crf_rslt_dir = osp.join(aff_rslt_dir, 'seed/best')\n",
    "\n",
    "# * 读取图片的原始CAM。\n",
    "loaded = np.load(osp.join(infer_rslt_dir, 'cam', f'{img_id}.npz'))\n",
    "\n",
    "fg_cls = loaded['fg_cls']\n",
    "\n",
    "cam_ori = resize_cam_cuda(torch.as_tensor(loaded['cam'], device=device), (ori_h, ori_w))\n",
    "score_ori = cam2score_cuda(cam_ori, (ori_h, ori_w), resize_first=True)\n",
    "\n",
    "viz_fg_cam(img, lb, cam_ori, 'cam_ori', with_bg=False)\n",
    "viz_fg_cam(img, lb, score_ori, 'score_ori', with_bg=False)\n",
    "\n",
    "# * 读取图片的affed CAM。\n",
    "loaded = np.load(osp.join(aff_rslt_dir, 'cam_affed', f'{img_id}.npz'))\n",
    "\n",
    "cam_affed = resize_cam_cuda(torch.as_tensor(loaded['cam'], device=device), (ori_h, ori_w))  # 有可能aff at score。\n",
    "score_affed = cam2score_cuda(cam_affed, (ori_h, ori_w), resize_first=True)\n",
    "\n",
    "viz_fg_cam(img, lb, cam_affed, 'cam_affed', with_bg=False)\n",
    "viz_fg_cam(img, lb, score_affed, 'score_affed', with_bg=False)\n",
    "\n",
    "# * 读取图片的CRF后score。\n",
    "loaded = np.load(osp.join(crf_rslt_dir, 'data', f'{img_id}.npz'))\n",
    "bg_fg_score = torch.as_tensor(loaded['bg_fg_score'], device=device)\n",
    "\n",
    "viz_fg_cam(img, lb, bg_fg_score, 'bg_fg_score', with_bg=True)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 提示法-V1"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 选取合适输入，处理为密集提示。"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 定义密集提示生成函数。"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def const_fg_bg_ign(bg_fg_score: torch.Tensor,\n",
    "                    thresh: float,\n",
    "                    fg_val:float, bg_val: float, ign_val: float=0, pad_val: float=-40.) -> np.ndarray:\n",
    "    bg_fg_score = resize_cam_cuda(bg_fg_score, (sam_low_res_input_h, sam_low_res_input_w))\n",
    "\n",
    "    seed = torch.argmax(bg_fg_score, dim=0)\n",
    "    low_res_logit_mask = torch.full_like(bg_fg_score[1:, :, :], bg_val)\n",
    "\n",
    "    for i in range(1, bg_fg_score.shape[0]):\n",
    "        conf_mask = bg_fg_score[i] > thresh\n",
    "        low_res_logit_mask[i-1, (seed == i) & conf_mask] = fg_val\n",
    "        low_res_logit_mask[i-1, (seed == i) & (~conf_mask)] = ign_val\n",
    "\n",
    "    h, w = low_res_logit_mask.shape[-2:]\n",
    "    pad_h = low_res_mask_size - h\n",
    "    pad_w = low_res_mask_size - w\n",
    "    low_res_logit_mask = F.pad(low_res_logit_mask, (0, pad_w, 0, pad_h), value=pad_val)  # 向右下填0，可以简化坐标变换。\n",
    "\n",
    "    return low_res_logit_mask"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 得到密集提示。"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# * 将score处理为SAM接收的形式（1x256x256，logit形式）。\n",
    "input_masks = const_fg_bg_ign(bg_fg_score,\n",
    "                              thresh=0.95,\n",
    "                              fg_val=4., bg_val=-10., ign_val=0., pad_val=-40.)\n",
    "viz_fg_cam(low_res_img, low_res_lb, input_masks, 'low_res_logit_mask', with_bg=False)\n",
    "input_masks = input_masks[:, None, :, :]"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 提取点提示"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 定义点提示生成函数。"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def merge_points_other_fg_as_bg(input_points: list[np.ndarray], input_labels: list[np.ndarray]) \\\n",
    "    -> tuple[list[np.ndarray], list[np.ndarray]]:\n",
    "    fg_points = [input_point[input_label.astype(bool)] for input_point, input_label in zip(input_points, input_labels)]\n",
    "    bg_points = [input_point[~input_label.astype(bool)] for input_point, input_label in zip(input_points, input_labels)]\n",
    "\n",
    "    merged_bg_points = []\n",
    "    for i in range(input_num := len(input_points)):\n",
    "        my_bg_point = bg_points[i]\n",
    "        others_fg_points = [fg_points[j] for j in range(input_num) if j != i]\n",
    "        merged_bg_points.append(np.concatenate([my_bg_point, *others_fg_points], axis=0))\n",
    "\n",
    "    merged_input_points, merged_input_labels = [], []\n",
    "    for fg_point, merged_bg_point in zip(fg_points, merged_bg_points):\n",
    "        merged_input_points.append(merged_input_point := np.concatenate([fg_point, merged_bg_point], axis=0))\n",
    "        merged_input_label = np.zeros(merged_input_point.shape[0], dtype=np.int64)\n",
    "        merged_input_label[:fg_point.shape[0]] = 1\n",
    "        merged_input_labels.append(merged_input_label)\n",
    "\n",
    "    return merged_input_points, merged_input_labels\n",
    "\n",
    "\n",
    "def is_ill_heat_map(heat_map: torch.Tensor,\n",
    "                    fg_min_area_rel: float=.05, bg_min_area_rel: float=.05,\n",
    "                    fg_thresh_rel: float=0., bg_thresh_rel: float=0.) -> bool | str:\n",
    "    # * 最大值小于0，最小值大于0，病态。\n",
    "    max_heat = heat_map.max()\n",
    "    min_heat = heat_map.min()\n",
    "    if max_heat < 0 or min_heat > 0:\n",
    "        return f\"max_heat: {max_heat}, min_heat: {min_heat}\"\n",
    "    # * 前景区域过小，病态。\n",
    "    if (fg_area := (heat_map > fg_thresh_rel * max_heat).sum()) < (fg_min_area := fg_min_area_rel * heat_map.numel()):\n",
    "        return f\"fg_area: {fg_area} < fg_min_area: {fg_min_area}\"\n",
    "    # * 背景区域过小，病态。\n",
    "    if (bg_area := (heat_map < bg_thresh_rel * min_heat).sum()) < (bg_min_area := bg_min_area_rel * heat_map.numel()):\n",
    "        return f\"bg_area: {bg_area} < bg_min_area: {bg_min_area}\"\n",
    "    return False\n",
    "\n",
    "\n",
    "def get_point_peak_local(heat_map: torch.Tensor,\n",
    "                         fg_min_area_rel: float=.05, bg_min_area_rel: float=.05,\n",
    "                         fg_thresh_rel: float=0., bg_thresh_rel: float=0.,\n",
    "                         min_distance_rel: float=None,\n",
    "                         fg_peak_thresh_rel=None, bg_peak_thresh_rel=None,\n",
    "                         exclude_border=True,\n",
    "                         max_fg_points_num=np.inf, max_bg_points_num=np.inf,\n",
    "                         footprint=None, p_norm=np.inf) -> tuple[np.ndarray, np.ndarray] | tuple[None, None]:\n",
    "    # * 检查是否病态。\n",
    "    if ill_msg := is_ill_heat_map(heat_map,\n",
    "                                  fg_min_area_rel, bg_min_area_rel,\n",
    "                                  fg_thresh_rel, bg_thresh_rel):\n",
    "        warnings.warn(f\"发现病态热力图: {ill_msg}\")\n",
    "        return None, None\n",
    "\n",
    "    # * 计算前背景共享参数。\n",
    "    min_distance = round(np.sqrt(heat_map.numel()) * min_distance_rel) if min_distance_rel else 1\n",
    "\n",
    "    # * 找前景点。\n",
    "    # ** 构造前景heat map。\n",
    "    max_heat = heat_map.max()\n",
    "    fg_heat_map = F.threshold(heat_map, fg_thresh_abs := fg_thresh_rel * max_heat, fg_thresh_abs)\n",
    "    # ** 找极大值作为前景点。\n",
    "    fg_points = peak_local_max(fg_heat_map.cpu().numpy(),\n",
    "                               min_distance=min_distance,\n",
    "                               threshold_abs=None, threshold_rel=fg_peak_thresh_rel,\n",
    "                               exclude_border=exclude_border, num_peaks=max_fg_points_num,\n",
    "                               footprint=footprint, p_norm=p_norm)[:, ::-1]\n",
    "\n",
    "    # * 找背景点。\n",
    "    # ** 构造前景heat map。\n",
    "    inv_heat_map = -heat_map\n",
    "    max_inv_heat = inv_heat_map.max()\n",
    "    bg_heat_map = F.threshold(inv_heat_map, bg_thresh_abs := bg_thresh_rel * max_inv_heat, bg_thresh_abs)\n",
    "    # ** 找极大值作为前景点。\n",
    "    bg_points = peak_local_max(bg_heat_map.cpu().numpy(),\n",
    "                               min_distance=min_distance,\n",
    "                               threshold_abs=None, threshold_rel=bg_peak_thresh_rel,\n",
    "                               exclude_border=exclude_border, num_peaks=max_bg_points_num,\n",
    "                               footprint=footprint, p_norm=p_norm)[:, ::-1]\n",
    "\n",
    "    # * 构造最终输出。\n",
    "    points = np.concatenate([fg_points, bg_points], axis=0)\n",
    "    labels = np.zeros(points.shape[0], dtype=np.int64)\n",
    "    labels[:fg_points.shape[0]] = 1\n",
    "\n",
    "    return points, labels\n",
    "\n",
    "\n",
    "def get_point_thresh_label_center(heat_map: torch.Tensor,\n",
    "                                  fg_min_area_rel: float=.05, bg_min_area_rel: float=.05,\n",
    "                                  fg_thresh_rel: float=0., bg_thresh_rel: float=0.,\n",
    "                                  fg_area_thresh_rel=.05, bg_area_thresh_rel=.05,\n",
    "                                  # exclude_border=True,\n",
    "                                  # max_fg_points_num=np.inf, max_bg_points_num=np.inf,\n",
    "                                  ) -> tuple[np.ndarray, np.ndarray] | tuple[None, None]:\n",
    "    # * 检查是否病态。\n",
    "    if ill_msg := is_ill_heat_map(heat_map,\n",
    "                                  fg_min_area_rel, bg_min_area_rel,\n",
    "                                  fg_thresh_rel, bg_thresh_rel):\n",
    "        warnings.warn(f\"发现病态热力图: {ill_msg}\")\n",
    "        return None, None\n",
    "\n",
    "    heat_map = heat_map.cpu().numpy()\n",
    "\n",
    "    # * 找前景点。\n",
    "    # ** 构造前景响应区域。\n",
    "    max_heat = heat_map.max()\n",
    "    fg_region = heat_map > fg_thresh_rel * max_heat\n",
    "    label_fg_region = label(fg_region, connectivity=2)\n",
    "    fg_regions = regionprops(label_fg_region)\n",
    "\n",
    "    # ** 找极大值作为前景点。\n",
    "    fg_points = []\n",
    "    for region in fg_regions:\n",
    "        if region.area < fg_area_thresh_rel * heat_map.size:\n",
    "            continue\n",
    "        fg_points.append(region.centroid[::-1])\n",
    "    fg_points = np.asarray(fg_points, dtype=np.int64)\n",
    "\n",
    "    # * 找背景点。\n",
    "    # ** 构造前景heat map。\n",
    "    inv_heat_map = -heat_map\n",
    "    max_inv_heat = inv_heat_map.max()\n",
    "    bg_region = inv_heat_map > bg_thresh_rel * max_inv_heat\n",
    "    label_bg_region = label(bg_region, connectivity=2)\n",
    "    bg_regions = regionprops(label_bg_region)\n",
    "    # ** 找极大值作为前景点。\n",
    "    bg_points = []\n",
    "    for region in bg_regions:\n",
    "        if region.area < bg_area_thresh_rel * heat_map.size:\n",
    "            continue\n",
    "        bg_points.append(region.centroid[::-1])\n",
    "    bg_points = np.asarray(bg_points, dtype=np.int64)\n",
    "\n",
    "    # * 构造最终输出。\n",
    "    points = np.concatenate([fg_points, bg_points], axis=0)\n",
    "    labels = np.zeros(points.shape[0], dtype=np.int64)\n",
    "    labels[:fg_points.shape[0]] = 1\n",
    "\n",
    "    return points, labels"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 获取点提示。"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# NOTE 可以尝试CAM或CAM affed作为点提示热度图先验。\n",
    "heat_maps = cam_ori\n",
    "# heat_maps = cam_affed\n",
    "\n",
    "input_points, input_labels = [], []\n",
    "for heat_map in heat_maps:\n",
    "    # input_point, input_label = get_point_peak_local(heat_map,\n",
    "    #                                                 fg_min_area_rel=.05, bg_min_area_rel=.05,\n",
    "    #                                                 fg_thresh_rel=0., bg_thresh_rel=0.,\n",
    "    #                                                 min_distance_rel=0.05,\n",
    "    #                                                 fg_peak_thresh_rel=0.2, bg_peak_thresh_rel=0.2,\n",
    "    #                                                 exclude_border=True,\n",
    "    #                                                 max_fg_points_num=10, max_bg_points_num=5,\n",
    "    #                                                 footprint=None, p_norm=np.inf)\n",
    "    input_point, input_label = get_point_thresh_label_center(heat_map,\n",
    "                                                             fg_min_area_rel=.05, bg_min_area_rel=.05,\n",
    "                                                             fg_thresh_rel=0.05, bg_thresh_rel=0.2,\n",
    "                                                             fg_area_thresh_rel=.001, bg_area_thresh_rel=.01)\n",
    "    input_points.append(input_point)\n",
    "    input_labels.append(input_label)\n",
    "viz_fg_cam(img, lb, heat_maps, '', with_bg=False, input_points=input_points, input_labels=input_labels)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "### 融合点提示。"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "input_points, input_labels = merge_points_other_fg_as_bg(input_points, input_labels)\n",
    "viz_fg_cam(img, lb, heat_maps, '', with_bg=False, input_points=input_points, input_labels=input_labels)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# SAM推理"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 输入SAM，获取分割结果"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# * 将score作为mask直接输入SAM，观察输出。\n",
    "predictor.set_image(img)\n",
    "\n",
    "sam_prompts_outs = []\n",
    "for input_point, input_label, input_mask in zip(input_points, input_labels, input_masks, strict=True):\n",
    "    spo = ADict()\n",
    "    spo.input_point = input_point\n",
    "    spo.input_label = input_label\n",
    "    spo.input_mask = None\n",
    "    sam_prompts_outs.append(spo)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "for spo in sam_prompts_outs:\n",
    "    spo.mask, spo.score, spo.low_res_logit_mask = predictor.predict(\n",
    "        point_coords=spo.input_point,\n",
    "        point_labels=spo.input_label,\n",
    "        mask_input = spo.input_mask,\n",
    "        multimask_output=False,\n",
    "        return_logits=False\n",
    "    )\n",
    "\n",
    "    spo.logit_mask, _, _ = predictor.predict(\n",
    "        point_coords=spo.input_point,\n",
    "        point_labels=spo.input_label,\n",
    "        mask_input = spo.input_mask,\n",
    "        multimask_output=False,\n",
    "        return_logits=True\n",
    "    )"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 打印SAM分割可视化"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# * 打印提示和分割结果。\n",
    "for i, spo in enumerate(sam_prompts_outs):\n",
    "    print(f\"=============== Plot spo {i}===============\")\n",
    "\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(img)\n",
    "    show_mask(spo.mask, plt.gca())\n",
    "    show_points(spo.input_point, spo.input_label, plt.gca())\n",
    "    plt.title(f\"Score: {spo.score[0]:.3f}\", fontsize=18)\n",
    "    plt.axis('on')\n",
    "    plt.show()\n",
    "\n",
    "    # * 打印low_res_logit_mask。\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(spo.low_res_logit_mask.transpose(1, 0, 2).reshape(low_res_mask_size, -1))\n",
    "    plt.colorbar(fraction=0.1)\n",
    "    plt.axis('on')\n",
    "    plt.show()\n",
    "\n",
    "    # * 打印logit_mask。\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(spo.logit_mask.transpose(1, 0, 2).reshape(img.shape[0], -1))\n",
    "    plt.colorbar(fraction=0.025)\n",
    "    plt.axis('on')\n",
    "    plt.show()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## 打印SAM分割数值分析"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "for i, spo in enumerate(sam_prompts_outs):\n",
    "    print(f\"=============== Stat spo {i}===============\")\n",
    "\n",
    "    print(f\"{spo.logit_mask.mean()=}, {spo.logit_mask.std()=}, {spo.logit_mask.min()=}, {spo.logit_mask.max()=}\")\n",
    "\n",
    "    for i, logit in enumerate(spo.low_res_logit_mask):\n",
    "        print(f\"mask {i} padding logit: {logit[-low_res_pad_h:, -low_res_pad_w:].mean()=}\")\n",
    "\n",
    "    for i, logit in enumerate(spo.logit_mask):\n",
    "        print(f\"mask {i} fg logit: {logit[spo.mask[i]].mean()=}\")\n",
    "        print(f\"mask {i} bg logit: {logit[~spo.mask[i]].mean()=}\")\n",
    "\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    for i, logit in enumerate(spo.low_res_logit_mask):\n",
    "        plt.subplot(1, 1, i + 1)\n",
    "        plt.hist(logit.flatten(), bins=50)\n",
    "        plt.yscale('log', base=10)\n",
    "        plt.tick_params(axis='both', which='major', labelsize=20)\n",
    "    plt.show()"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# 终点"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
